import numpy as np
import cv2
import mediapipe as mp

#get raw datas for one class
def get_class_gestures(class_name,data_num):
    gestures=[]
    cur_num=0
    # 初始化 Mediapipe 手部模块
    mp_hands = mp.solutions.hands
    hands = mp_hands.Hands()
    # 初始化绘图工具
    mp_drawing = mp.solutions.drawing_utils
    # 打开摄像头
    cap = cv2.VideoCapture(0)
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            continue
        # 将 BGR 图像转换为 RGB
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        # 进行手部关节检测
        results = hands.process(rgb_frame)
        # 如果检测到手部关节
        
        if results.multi_hand_landmarks:
            for landmarks in results.multi_hand_landmarks:
                # 在图像上绘制手部关节
                mp_drawing.draw_landmarks(frame, landmarks, mp_hands.HAND_CONNECTIONS)
                gesture=[[landmark.x,landmark.y,landmark.z] for landmark in landmarks.landmark]
                if len(gesture)==21:
                    gestures.append(gesture)
                    cur_num+=1
                    print("cur num: ",cur_num)
        if cur_num==data_num:
            break
        # 显示帧
        cv2.imshow('Hand Tracking', frame)
        # 按 'q' 键退出循环
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    # 释放资源
    cap.release()
    cv2.destroyAllWindows()
    return gestures,class_name
#get the neighbor matrix A
def save_A():
    A=np.zeros((21,21))
    A[0][1],A[1][2],A[2][3],A[3][4],A[0][5]=1,1,1,1,1
    A[5][6],A[6][7],A[7][8],A[9][10],A[10][11]=1,1,1,1,1
    A[11][12],A[13][14],A[14][15],A[15][16],A[17][18]=1,1,1,1,1
    A[18][19],A[19][20]=1,1
    A=A.astype(np.float32)
    np.save('./A.npy',A)

#input data,convert to list of tuple,each tuple is a channel of a node,and the index stands for node id
def save_data_to_graph(gestures,class_name):
    data=np.array(gestures).astype(np.float32)
    print(data.shape)
    np.save('./nodes/'+class_name+'.npy',data)

save_A()
# gestures,class_name=get_class_gestures("open_right",500)
# save_data_to_graph(gestures,class_name)


